import numpy as np
import pandas as pd
import re
import matplotlib.pyplot as plt


class XVG:
    M = re.compile(r'.*\"(.*)\".*')

    def __init__(self, xvg):
        self.xvg = xvg
        with open(xvg) as f:
            lines = f.readlines()
            paras = [l for l in lines if l.startswith('@')]
            data = [[float(i) for i in l.split()] for l in lines if not l.startswith('#') and not l.startswith('@')]
        self.title = self.parse_title(paras)
        self.xlabel, self.ylabel = self.parse_labels(paras)
        self.legends = self.parse_legends(paras)
        if len(self.legends) == 0:
            self.xvgdata = pd.DataFrame([d[1:] for d in data], columns=[self.ylabel], index=[d[0] for d in data])
        else:
            self.xvgdata = pd.DataFrame([d[1:] for d in data], columns=self.legends, index=[d[0] for d in data])

    def get_title(self):
        return self.title

    def get_labels(self):
        return self.xlabel, self.ylabel

    def get_legends(self):
        return self.legends

    def get_data(self):
        return self.xvgdata

    def parse_title(self, paras):
        title = re.match(self.M, [p for p in paras if p.split()[1] == 'title'][0]).group(1)
        return title

    def parse_labels(self, paras):
        xlabel = re.match(self.M, [p for p in paras if p.split()[1] == 'xaxis'][0]).group(1)
        ylabel = re.match(self.M, [p for p in paras if p.split()[1] == 'yaxis'][0]).group(1)
        xlabel = '$' + xlabel.replace(r'\S', '^{').replace(r'\N', '}').replace(' ', '\ ') + '$'
        ylabel = '$' + ylabel.replace(r'\S', '^{').replace(r'\N', '}').replace(' ', '\ ') + '$'
        return xlabel, ylabel

    def parse_legends(self, paras):
        legends = [re.match(self.M, s_line).group(1) for s_line in
                   [p for p in paras if p.startswith('@ s') and p.split()[2] == 'legend']]
        return legends

    def expand(self, new_xvg_list=[], join='outer'):
        if not isinstance(new_xvg_list, list):
            raise 'new_xvg_list must be list.'
        if len(new_xvg_list) == 0:
            raise 'No new xvg got.'
        for new_xvg in new_xvg_list:
            self.xvgdata = pd.concat([self.xvgdata, new_xvg.xvgdata], axis=1, join=join)
            self.legends = self.xvgdata.columns

    def to_csv(self, csv_name):
        self.xvgdata.to_csv(csv_name, index=True, header=self.title)

    def to_fig(self, fig_name):
        self.xvgdata.plot(xlabel=self.xlabel, ylabel=self.ylabel)
        plt.savefig(fig_name)
        plt.clf()


def xvg2csv(xvgname, csv_name):
    xvgfile = XVG(xvgname)
    xvgfile.get_data().to_csv(csv_name, index=False, header=xvgfile.get_title())


if __name__ == "__main__":
    xvg = XVG('test/xvg/TA_M0.xvg')
    xvg_20 = XVG('test/xvg/TA_M20.xvg')
    xvg_100 = XVG('test/xvg/TA_M20.xvg')
    xvg.expand([xvg_20, xvg_100])
    xvg.to_csv('test/TA_M0_M20_M100.csv')
